import os
import logging
from typing import List, Dict, Any, Optional
import numpy as np
import requests
from langchain.schema import Document
from langchain_community.vectorstores import Chroma
from config import Config
from config.thresholds import EmbeddingThresholds

logger = logging.getLogger(__name__)

class ArkEmbeddings:
    """方舟嵌入模型包装器 - 使用requests直接调用API"""
    
    def __init__(self, api_key: str, model_name: str = "doubao-embedding-large-text-250515"):
        self.api_key = api_key
        self.model_name = model_name
        self.api_base = "https://ark.cn-beijing.volces.com/api/v3"
        logger.info(f"初始化方舟嵌入模型: {self.model_name}")
    
    def encode(self, inputs: List[str], is_query: bool = False, mrl_dim: Optional[int] = None):
        """编码文本为嵌入向量"""
        try:
            if is_query:
                # 为查询添加指令以获得最佳性能
                inputs = [
                    f"Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: {i}"
                    for i in inputs
                ]
            
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }
            
            data = {
                "model": self.model_name,
                "input": inputs,
                "encoding_format": "float"
            }
            
            response = requests.post(
                f"{self.api_base}/embeddings",
                headers=headers,
                json=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                embeddings = [item["embedding"] for item in result["data"]]
                
                # 转换为numpy数组
                embedding = np.array(embeddings, dtype=np.float32)
                
                if mrl_dim is not None:
                    assert mrl_dim in [2048, 1024, 512, 256]
                    embedding = embedding[:, :mrl_dim]
                
                # 归一化以计算余弦相似度
                norms = np.linalg.norm(embedding, axis=1, keepdims=True)
                norms = np.where(norms == 0, 1, norms)  # 避免除零
                embedding = embedding / norms
                
                return embedding
            else:
                logger.error(f"嵌入API调用失败: {response.status_code} - {response.text}")
                raise Exception(f"API调用失败: {response.status_code}")
                
        except Exception as e:
            logger.error(f"嵌入编码失败: {e}")
            raise
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """为文档生成嵌入向量"""
        try:
            embeddings = self.encode(texts, is_query=False, mrl_dim=1024)
            return embeddings.tolist()
        except Exception as e:
            logger.error(f"文档嵌入失败: {e}")
            raise
    
    def embed_query(self, text: str) -> List[float]:
        """为查询生成嵌入向量"""
        try:
            embeddings = self.encode([text], is_query=True, mrl_dim=1024)
            return embeddings[0].tolist()
        except Exception as e:
            logger.error(f"查询嵌入失败: {e}")
            raise

class EmbeddingManager:
    """词嵌入管理器"""
    
    def __init__(self):
        # 初始化方舟嵌入模型
        self.embeddings = ArkEmbeddings(
            api_key=Config.OPENAI_API_KEY,
            model_name=Config.EMBEDDING_MODEL_NAME
        )
        logger.info(f"初始化词嵌入模型: {Config.EMBEDDING_MODEL_NAME}")
    
    def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """为文本列表生成词嵌入"""
        try:
            logger.info(f"为 {len(texts)} 个文本生成词嵌入")
            
            # 批量生成词嵌入
            embeddings = self.embeddings.embed_documents(texts)
            
            logger.info(f"词嵌入生成完成，维度: {len(embeddings[0]) if embeddings else 0}")
            return embeddings
            
        except Exception as e:
            logger.error(f"生成词嵌入失败: {str(e)}")
            raise
    
    def generate_single_embedding(self, text: str) -> List[float]:
        """为单个文本生成词嵌入"""
        try:
            logger.info(f"为文本生成词嵌入: {text[:50]}...")
            
            embedding = self.embeddings.embed_query(text)
            
            logger.info(f"词嵌入生成完成，维度: {len(embedding)}")
            return embedding
            
        except Exception as e:
            logger.error(f"生成词嵌入失败: {str(e)}")
            raise
    
    def create_vector_store(
        self,
        documents: List[Document],
        collection_name: str,
        persist_directory: str
    ) -> Chroma:
        """创建向量存储"""
        try:
            logger.info(f"创建向量存储，集合: {collection_name}")
            
            # 这里会自动进行词嵌入
            vector_store = Chroma.from_documents(
                documents=documents,
                embedding=self.embeddings,  # 自动为每个文档生成词嵌入
                persist_directory=persist_directory,
                collection_name=collection_name
            )
            
            # 持久化存储
            vector_store.persist()
            
            logger.info(f"向量存储创建完成，文档数: {len(documents)}")
            return vector_store
            
        except Exception as e:
            logger.error(f"创建向量存储失败: {str(e)}")
            raise
    
    def add_documents_to_store(
        self,
        documents: List[Document],
        collection_name: str,
        persist_directory: str
    ) -> Chroma:
        """向现有向量存储添加文档"""
        try:
            logger.info(f"向向量存储添加文档，集合: {collection_name}")
            
            # 加载现有向量存储
            vector_store = Chroma(
                persist_directory=persist_directory,
                embedding_function=self.embeddings,
                collection_name=collection_name
            )
            
            # 添加新文档（自动生成词嵌入）
            vector_store.add_documents(documents)
            
            # 持久化
            vector_store.persist()
            
            logger.info(f"文档添加完成，新增文档数: {len(documents)}")
            return vector_store
            
        except Exception as e:
            logger.error(f"添加文档失败: {str(e)}")
            raise
    
    def similarity_search(
        self,
        query: str,
        collection_name: str,
        
        persist_directory: str,
        top_k: int = 5
    ) -> List[Document]:
        """相似性搜索"""
        try:
            logger.info(f"执行相似性搜索: {query[:50]}...")
            
            # 加载向量存储
            vector_store = Chroma(
                persist_directory=persist_directory,
                embedding_function=self.embeddings,
                collection_name=collection_name
            )
            
            # 执行搜索（查询会自动生成词嵌入）
            results = vector_store.similarity_search(query, k=top_k)
            
            logger.info(f"搜索完成，找到 {len(results)} 个相关文档")
            return results
            
        except Exception as e:
            logger.error(f"相似性搜索失败: {str(e)}")
            raise
    
    def get_embedding_info(self) -> Dict[str, Any]:
        """获取词嵌入模型信息"""
        return {
            "model_name": Config.EMBEDDING_MODEL_NAME,
            "api_base": Config.OPENAI_BASE_URL,
            "embedding_dimension": EmbeddingThresholds.EMBEDDING_DIMENSION,
            "max_tokens": EmbeddingThresholds.MAX_TEXT_LENGTH
        } 